package service

import (
	"strconv"
	"sugar/controller"
	"sugar/errorcode"
	"sugar/utils"

	"github.com/gin-gonic/gin"
	"gopkg.in/guregu/null.v3"
)

type createTrainJobParams struct {
	Name           null.String `json:"name"`
	Version        null.Int    `json:"version"`
	ProjectID      null.Int    `json:"project_id"`
	Type           null.Int    `json:"type"`
	Description    null.String `json:"description"`
	Budget         null.String `json:"budget"`
	TrainDatasetID null.Int    `json:"train_dataset_id"`
	ValDatasetID   null.Int    `json:"val_dataset_id"`
	ModelID        null.Int    `json:"model_id"`
	ModelCommitID  null.String `json:"model_commit_id"`
	ModelArgs      null.String `json:"model_args"`
	TrainArgs      null.String `json:"train_args"`
	UseAutoml      int         `json:"use_automl"`
	DestScore      string      `json:"dest_score"`
	MaxTrials      int         `json:"maxtrials"`
}

type getTrainJobParams struct {
	ID null.Int `json:"id"`
}

type updateTrainJobStatusParams struct {
	PendingReplicas int `json:"pendingreplicas"`

	RunningReplicas int `json:"runningreplicas"`

	ErrorReplicas int `json:"errorreplicas"`

	CompleteReplicas int `json:"completereplicas"`

	Status string `json:"status"`
}

// @Summary 创建训练任务
// @Description 创建训练任务
// @Tags 训练任务管理
// @Accept  json
// @Produce json
// @Param   params     body    createTrainJobParams     true        "创建训练任务参数"
// @Success 200 {object}  response.Response	"Success data为null"
// @Failure 500  {object}  response.Response  "内部错误 data为null"
// @Failure 401  {object}  response.Response  "token错误 data为null"
// @Failure 404  {object}  response.Response  "无法访问 data为null"
// @Router /trainjob [post]
func createTrainJob(c *gin.Context, tokenMap map[string]interface{}) {
	paramOk := utils.CheckUUID4MustExist(tokenMap["user_id"])
	if !paramOk {
		panic(errorcode.ParamError)
	}
	jsonParams := createTrainJobParams{
		UseAutoml: 0,
	}
	err := c.BindJSON(&jsonParams)
	if err != nil {
		panic(errorcode.ParamError)
	}
	paramOk = utils.CheckStringMustExist(jsonParams.Name) &&
		utils.CheckIntMustExist(jsonParams.TrainDatasetID, jsonParams.ValDatasetID, jsonParams.ModelID, jsonParams.Type) &&
		utils.CheckIntIfExist(jsonParams.Version) &&
		utils.CheckStringIfExist(jsonParams.Description, jsonParams.Budget, jsonParams.ModelArgs, jsonParams.TrainArgs, jsonParams.ModelCommitID)
	if !paramOk {
		panic(errorcode.ParamError)
	}

	trainJobOp := controller.NewOp()
	res := trainJobOp.CreateTrainJob(tokenMap["user_id"].(string), jsonParams.Name.String, int(jsonParams.Version.Int64), int(jsonParams.Type.Int64), jsonParams.ProjectID,
		jsonParams.Description, jsonParams.Budget, jsonParams.TrainDatasetID.Int64, jsonParams.ValDatasetID.Int64, jsonParams.ModelID.Int64,
		jsonParams.ModelCommitID, jsonParams.ModelArgs, jsonParams.TrainArgs, jsonParams.DestScore, jsonParams.MaxTrials, jsonParams.UseAutoml)
	c.JSON(res.StatusCode, res)
}

// @Summary 停止训练任务
// @Description 停止训练任务
// @Tags 训练任务管理
// @Accept  json
// @Produce json
// @Param   params     body    getTrainJobParams     true        "训练任务ID"
// @Success 200 {object}  response.Response	"Success data为null"
// @Failure 500  {object}  response.Response  "内部错误 data为null"
// @Failure 401  {object}  response.Response  "token错误 data为null"
// @Failure 404  {object}  response.Response  "无法访问 data为null"
// @Router /trainjob/stop [post]
func stopTrainJob(c *gin.Context, tokenMap map[string]interface{}) {
	paramOk := utils.CheckUUID4MustExist(tokenMap["user_id"])
	if !paramOk {
		panic(errorcode.ParamError)
	}
	var jsonParams getTrainJobParams
	err := c.BindJSON(&jsonParams)
	if err != nil {
		panic(errorcode.ParamError)
	}
	paramOk = utils.CheckIntMustExist(jsonParams.ID)
	if !paramOk {
		panic(errorcode.ParamError)
	}
	trainJobOp := controller.NewOp()
	res := trainJobOp.StopTrainJob(tokenMap["user_id"].(string), jsonParams.ID.Int64)
	c.JSON(res.StatusCode, res)
}

// @Summary 获取训练任务详情
// @Description 获取训练任务详情
// @Tags 训练任务管理
// @Accept  json
// @Produce json
// @Param   params     body    getTrainJobParams     true        "训练任务ID"
// @Success 200 {object}  response.Response	"Success data为null"
// @Failure 500  {object}  response.Response  "内部错误 data为null"
// @Failure 401  {object}  response.Response  "token错误 data为null"
// @Failure 404  {object}  response.Response  "无法访问 data为null"
// @Router /trainjob [get]
func getTrainJob(c *gin.Context, tokenMap map[string]interface{}) {
	paramOk := utils.CheckUUID4MustExist(tokenMap["user_id"])
	if !paramOk {
		panic(errorcode.ParamError)
	}
	var jsonParams getTrainJobParams
	err := c.BindQuery(&jsonParams)
	if err != nil {
		panic(errorcode.ParamError)
	}
	paramOk = utils.CheckIntMustExist(jsonParams.ID)
	if !paramOk {
		panic(errorcode.ParamError)
	}
	trainJobOp := controller.NewOp()
	res := trainJobOp.GetTrainJob(tokenMap["user_id"].(string), jsonParams.ID.Int64)
	c.JSON(res.StatusCode, res)
}

// @Summary 获取训练任务列表
// @Description 获取训练任务列表
// @Tags 训练任务管理
// @Accept  json
// @Produce json
// @Success 200 {object}  response.Response	"Success data为null"
// @Failure 500  {object}  response.Response  "内部错误 data为null"
// @Failure 401  {object}  response.Response  "token错误 data为null"
// @Failure 404  {object}  response.Response  "无法访问 data为null"
// @Router /trainjob/list [get]
func listTrainJob(c *gin.Context, tokenMap map[string]interface{}) {
	paramOk := utils.CheckUUID4MustExist(tokenMap["user_id"])
	if !paramOk {
		panic(errorcode.ParamError)
	}
	trainJobOp := controller.NewOp()
	res := trainJobOp.ListTrainJob(tokenMap["user_id"].(string))
	c.JSON(res.StatusCode, res)
}

// @Summary 更新训练任务状态
// @Description 更新训练任务状态
// @Tags 训练任务管理
// @Accept  json
// @Produce json
// @Param   params     body    updateTrainJobStatusParams     true        "训练任务状态"
// @Success 200 {object}  response.Response	"Success data为null"
// @Failure 500  {object}  response.Response  "内部错误 data为null"
// @Failure 401  {object}  response.Response  "token错误 data为null"
// @Failure 404  {object}  response.Response  "无法访问 data为null"
// @Router /trainjob/status/:id [post]
func updateTrainJobStatus(c *gin.Context) {
	// paramOk := utils.CheckUUID4MustExist(tokenMap["user_id"])
	// if !paramOk {
	// 	panic(errorcode.ParamError)
	// }
	var jsonParams updateTrainJobStatusParams
	err := c.BindJSON(&jsonParams)
	if err != nil {
		panic(errorcode.ParamError)
	}
	paramOk := utils.CheckStringMustExist(jsonParams.Status)
	if !paramOk {
		panic(errorcode.ParamError)
	}
	id, bexist := c.Params.Get("id")
	if !bexist {
		panic(errorcode.ParamError)
	}
	nid, _ := strconv.Atoi(id)
	trainJobOp := controller.NewOp()
	res := trainJobOp.UpdateTrainJobStatus(int64(nid), jsonParams.Status)
	c.JSON(res.StatusCode, res)
}
